package com.askey.bitratedemo;

import android.media.MediaCodec;
import android.media.MediaCodecInfo;
import android.media.MediaCodecList;
import android.media.MediaFormat;
import android.util.Log;
import android.view.Surface;

import java.io.IOException;
import java.nio.ByteBuffer;

import static android.media.MediaCodec.BUFFER_FLAG_CODEC_CONFIG;

public class VideoEncoder {
    private static final String TAG = " VideoEnc ";
    private EncoderConfig mConfig;
    private MediaCodec mMediaCodec;
    private MediaCodec.Callback mCallback;
    private OnMediaFormatReadyListener mMediaFormatReadyListener;

    public static class EncoderConfig {
        public String type;
        public int width;
        public int height;
        public int bitrate;
        public int framerate;
        public int mode;

        public String toString() {
            return width + " x " + height + "  " + ((mode == 0) ? "CQ" : (mode == 1) ? "VBR" : "CBR") + "  " + bitrate / 1024 + "kb/s"
                    ;
        }
    }

    public VideoEncoder(EncoderConfig config, OnMediaFormatReadyListener listener) {
        mConfig = config;
        mMediaFormatReadyListener = listener;
    }

    public interface OnMediaFormatReadyListener {
        void onMediaFormatReady(MediaFormat format);
    }

    // Call me after prepare.
    public Surface getSurface() {
        return mSurface;
    }

    private Surface mSurface;

    public void prepare() {
        final MediaFormat format = MediaFormat.createVideoFormat(mConfig.type, mConfig.width, mConfig.height);
        format.setInteger(MediaFormat.KEY_COLOR_FORMAT, MediaCodecInfo.CodecCapabilities.COLOR_FormatSurface);  // API >= 18
        format.setInteger(MediaFormat.KEY_BIT_RATE, mConfig.bitrate);
        format.setInteger(MediaFormat.KEY_FRAME_RATE, mConfig.framerate);
        format.setInteger(MediaFormat.KEY_I_FRAME_INTERVAL, 1);
        format.setInteger(MediaFormat.KEY_BITRATE_MODE, mConfig.mode);
        Log.d(TAG, "format: " + format);

        try {
            mMediaCodec = MediaCodec.createEncoderByType(mConfig.type);
        } catch (IOException e) {
            Log.e(TAG, "" + e.toString());
            return;
        }

        if (mCallback == null) mCallback = callback;
        mMediaCodec.setCallback(mCallback);
        mMediaCodec.configure(format, null, null, MediaCodec.CONFIGURE_FLAG_ENCODE);
        mSurface = mMediaCodec.createInputSurface();    // API >= 18
        mMediaCodec.start();
    }


    public void release() {
        if (mSurface != null) mSurface.release();

        mMediaCodec.stop();
        mMediaFormatReadyListener = null;
        mMediaCodec.release();
    }

    public interface OutputDataCallback {
        //Warning: ByteBuffer will be released once the function returned.
        void onDataAvailable(ByteBuffer bb, MediaCodec.BufferInfo info);
    }

    private Object mListenerSync = new Object();
    private volatile OutputDataCallback mOutputDataListener;

    public void setOutputDataListener(OutputDataCallback lintener) {
        synchronized (mListenerSync) {
            mOutputDataListener = lintener;
        }
    }

    public MediaCodec.Callback callback = new MediaCodec.Callback() {

        @Override
        public void onInputBufferAvailable(MediaCodec codec, int index) {

        }

        @Override
        public void onOutputBufferAvailable(MediaCodec codec, int index, MediaCodec.BufferInfo info) {
            if (mOutputDataListener != null) {
                if ((info.flags & BUFFER_FLAG_CODEC_CONFIG) != 0) {
                    Log.d(TAG, "Skip BUFFER_FLAG_CODEC_CONFIG.");
                    return;
                }
                ByteBuffer outputBuffer = codec.getOutputBuffer(index);
                mOutputDataListener.onDataAvailable(outputBuffer, info);
            }
            codec.releaseOutputBuffer(index, false);
        }

        @Override
        public void onError(MediaCodec codec, MediaCodec.CodecException e) {
            Log.e(TAG, "Codec error.");
        }

        @Override
        public void onOutputFormatChanged(MediaCodec codec, MediaFormat format) {
            if (mMediaFormatReadyListener != null) {
                mMediaFormatReadyListener.onMediaFormatReady(format);
            }
        }
    };

    protected final MediaCodecInfo selectVideoCodec(final String mimeType) {
        Log.d(TAG, "selectVideoCodec:");

        // get the list of available codecs
        MediaCodecList codecList = new MediaCodecList(MediaCodecList.ALL_CODECS);
        MediaCodecInfo[] codecInfos = codecList.getCodecInfos();
        for (MediaCodecInfo codecInfo : codecInfos) {

            if (!codecInfo.isEncoder()) {   // skipp decoder
                continue;
            }
            // select first codec that match a specific MIME type and color format
            final String[] types = codecInfo.getSupportedTypes();
            for (int j = 0; j < types.length; j++) {
                if (types[j].equalsIgnoreCase(mimeType)) {
                    Log.d(TAG, "codec:" + codecInfo.getName() + ",MIME=" + types[j]);
                    final int format = selectColorFormat(codecInfo, mimeType);
                    if (format > 0) {
                        return codecInfo;
                    }
                }
            }
        }
        return null;
    }

    protected final int selectColorFormat(final MediaCodecInfo codecInfo, final String mimeType) {
        Log.d(TAG, "selectColorFormat: ");
        int result = 0;
        final MediaCodecInfo.CodecCapabilities caps;
        try {
            Thread.currentThread().setPriority(Thread.MAX_PRIORITY);
            caps = codecInfo.getCapabilitiesForType(mimeType);
        } finally {
            Thread.currentThread().setPriority(Thread.NORM_PRIORITY);
        }
        int colorFormat;
        for (int i = 0; i < caps.colorFormats.length; i++) {
            colorFormat = caps.colorFormats[i];
            if (isRecognizedVideoFormat(colorFormat)) {
                if (result == 0)
                    result = colorFormat;
                break;
            }
        }
        if (result == 0)
            Log.e(TAG, "couldn't find a good color format for " + codecInfo.getName() + " / " + mimeType);
        return result;
    }

    private final boolean isRecognizedVideoFormat(final int colorFormat) {
        final int n = recognizedFormats != null ? recognizedFormats.length : 0;
        for (int i = 0; i < n; i++) {
            if (recognizedFormats[i] == colorFormat) {
                return true;
            }
        }
        return false;
    }

    protected static int[] recognizedFormats;

    static {
        recognizedFormats = new int[]{
//          MediaCodecInfo.CodecCapabilities.COLOR_FormatYUV420Planar,
//          MediaCodecInfo.CodecCapabilities.COLOR_FormatYUV420SemiPlanar,
//          MediaCodecInfo.CodecCapabilities.COLOR_QCOM_FormatYUV420SemiPlanar,
                MediaCodecInfo.CodecCapabilities.COLOR_FormatSurface,
        };
    }

    private volatile MediaCodec.BufferInfo mBufferInfo = new MediaCodec.BufferInfo();
    private volatile long mLastPts;

    public void setCallback(MediaCodec.Callback muxer) {
        mCallback = muxer;
    }
}
